# encoding=utf-8
import os

import torch
import torch.nn.functional as F
from model import PromptBert
from transformers import BertTokenizer
import os

Bert_path = '/data0/jianyu10/PTM/huggingface_model_cache/chinese-roberta-wwm-ext'

tokenizer = BertTokenizer.from_pretrained("hfl/chinese-roberta-wwm-ext", cache_dir=Bert_path)
DEVICE = torch.device('cpu' if torch.cuda.is_available() else 'cuda')


class ptConfig:
    root_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
    model_path = '/data0/jianyu10/PTM/huggingface_model_cache/chinese-roberta-wwm-ext'
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    dropout_prob = 0.25
    mask_ids = 103
    dropout = 0.15
    save_path = root_path + '/UseModel/'
    onnxpath = root_path + '/UseModel/' + 'prompt.onnx'


class PtModel:
    def __init__(self, config):
        self.model = PromptBert(config)
        self.model.to(DEVICE)
        path = 'UseModel/train_loss_best.pth.tar'
        chechpoint = torch.load(path)
        self.model.load_state_dict(chechpoint['state_dict'])
        self.model.eval()

    def gettmp(self, title):
        sentence = f'{title}，它的意思是[MASK]。'
        sen_tmp = f'{title}，这句话的意思是[MASK]。'
        return sentence, sen_tmp

    def inference(self, orititle, simtitle):
        with torch.no_grad():
            ori, oritmp = self.gettmp(orititle)
            pos, postmp = self.gettmp(simtitle)
            oriidx = tokenizer.encode_plus(ori, return_tensors='pt').get('input_ids')
            posidx = tokenizer.encode_plus(pos, return_tensors='pt').get('input_ids')
            oritmpidx = tokenizer.encode_plus(oritmp, return_tensors='pt').get('input_ids')
            postmpidx = tokenizer.encode_plus(postmp, return_tensors='pt').get('input_ids')
            oriten = self.model(oriidx, oritmpidx).view(1,-1)
            posten = self.model(posidx, postmpidx).view(1,-1)
            sim = F.cosine_similarity(oriten, posten, dim=-1)
            return {'orititle': oriten.numpy(), 'simtitle': posten.numpy(), 'simscore': sim.item()}


if __name__ == '__main__':
    pass
